#!/usr/bin/env python3
"""
视觉学习提示生成器

整合周期性、质量感知和记忆掩码提示的统一框架
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Dict, List, Tuple, Optional, Any, Union

import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parent.parent.parent))
from .base_prompt import BasePromptGenerator
from .periodic_prompt import PeriodicPromptGenerator
from .quality_prompt import QualityAwarePromptGenerator
from .memory_prompt import MemoryMaskPromptGenerator
from .self_prompt import SelfPromptGenerator
from utils.logger import get_logger


class VisualLearningPromptGenerator(BasePromptGenerator):
    """视觉学习提示生成器 - 统一框架"""
    
    def __init__(self, 
                 input_dim: int = 256,
                 prompt_dim: int = 256,
                 enable_periodic: bool = True,
                 enable_quality: bool = True,
                 enable_memory: bool = True,
                 enable_self_prompt: bool = True,
                 fusion_method: str = "attention",  # "attention", "concat", "weighted"
                 device: str = "cuda"):
        """
        初始化视觉学习提示生成器
        
        Args:
            input_dim: 输入特征维度
            prompt_dim: 提示特征维度
            enable_periodic: 启用周期性提示
            enable_quality: 启用质量感知提示
            enable_memory: 启用记忆掩码提示
            enable_self_prompt: 启用自提示机制
            fusion_method: 融合方法
            device: 设备
        """
        super().__init__(input_dim, prompt_dim, device)
        
        self.enable_periodic = enable_periodic
        self.enable_quality = enable_quality
        self.enable_memory = enable_memory
        self.enable_self_prompt = enable_self_prompt
        self.fusion_method = fusion_method
        
        # 初始化各个子模块
        self.modules = nn.ModuleDict()
        
        if enable_periodic:
            self.modules['periodic'] = PeriodicPromptGenerator(
                input_dim=input_dim,
                prompt_dim=prompt_dim,
                device=device
            )
        
        if enable_quality:
            self.modules['quality'] = QualityAwarePromptGenerator(
                input_dim=input_dim,
                prompt_dim=prompt_dim,
                device=device
            )
        
        if enable_memory:
            self.modules['memory'] = MemoryMaskPromptGenerator(
                input_dim=input_dim,
                prompt_dim=prompt_dim,
                device=device
            )
        
        if enable_self_prompt:
            self.modules['self_prompt'] = SelfPromptGenerator(
                input_dim=input_dim,
                prompt_dim=prompt_dim,
                device=device
            )
        
        # 融合网络
        if fusion_method == "attention":
            self.fusion_attention = nn.MultiheadAttention(
                embed_dim=prompt_dim,
                num_heads=8,
                dropout=0.1,
                batch_first=True
            )
        elif fusion_method == "weighted":
            num_modules = sum([enable_periodic, enable_quality, enable_memory, enable_self_prompt])
            self.fusion_weights = nn.Parameter(torch.ones(num_modules) / num_modules)
        
        # 最终提示生成器
        self.final_generator = nn.ModuleDict({
            'points': nn.Sequential(
                nn.Linear(prompt_dim, prompt_dim // 2),
                nn.ReLU(),
                nn.Linear(prompt_dim // 2, 2)
            ),
            'boxes': nn.Sequential(
                nn.Linear(prompt_dim, prompt_dim // 2),
                nn.ReLU(),
                nn.Linear(prompt_dim // 2, 4)
            ),
            'masks': nn.Sequential(
                nn.Linear(prompt_dim, prompt_dim // 2),
                nn.ReLU(),
                nn.Linear(prompt_dim // 2, 32 * 32)
            )
        })
        
        # 置信度预测器
        self.confidence_predictor = nn.Sequential(
            nn.Linear(prompt_dim, prompt_dim // 2),
            nn.ReLU(),
            nn.Linear(prompt_dim // 2, 1),
            nn.Sigmoid()
        )
        
        self.logger.info(f"视觉学习提示生成器初始化完成: {self.get_model_info()}")
    
    def generate_individual_prompts(self, 
                                  image: torch.Tensor,
                                  context: Optional[Dict[str, Any]] = None) -> Dict[str, Dict[str, torch.Tensor]]:
        """
        生成各个模块的提示
        
        Args:
            image: 输入图像 (B, C, H, W)
            context: 上下文信息
            
        Returns:
            各模块的提示字典
        """
        individual_prompts = {}
        
        # 周期性提示
        if self.enable_periodic:
            periodic_context = context.get('periodic', {}) if context else {}
            individual_prompts['periodic'] = self.modules['periodic'].generate_prompts(
                image, periodic_context
            )
        
        # 质量感知提示
        if self.enable_quality:
            quality_context = context.get('quality', {}) if context else {}
            individual_prompts['quality'] = self.modules['quality'].generate_prompts(
                image, quality_context
            )
        
        # 记忆掩码提示
        if self.enable_memory:
            memory_context = context.get('memory', {}) if context else {}
            individual_prompts['memory'] = self.modules['memory'].generate_prompts(
                image, memory_context
            )
        
        # 自提示机制
        if self.enable_self_prompt:
            self_prompt_context = context.get('self_prompt', {}) if context else {}
            individual_prompts['self_prompt'] = self.modules['self_prompt'].generate_prompts(
                image, self_prompt_context
            )
        
        return individual_prompts
    
    def fuse_prompts(self, 
                    individual_prompts: Dict[str, Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
        """
        融合各个模块的提示
        
        Args:
            individual_prompts: 各模块的提示字典
            
        Returns:
            融合后的提示字典
        """
        fused_prompts = {}
        
        # 收集所有模块的特征
        all_features = []
        module_names = []
        
        for module_name, prompts in individual_prompts.items():
            if 'attended_features' in prompts:
                features = prompts['attended_features']
            elif 'phase_info' in prompts:
                features = prompts['phase_info']
            elif 'memory_features' in prompts:
                features = prompts['memory_features']
            else:
                # 使用默认特征
                features = torch.randn(prompts['points'].size(0), self.prompt_dim, device=self.device)
            
            all_features.append(features)
            module_names.append(module_name)
        
        if not all_features:
            # 没有可用的模块，返回空提示
            batch_size = 1
            return {
                'points': torch.zeros(batch_size, 1, 2, device=self.device),
                'boxes': torch.zeros(batch_size, 1, 4, device=self.device),
                'masks': torch.zeros(batch_size, 1, 32, 32, device=self.device),
                'confidence': torch.zeros(batch_size, 1, device=self.device)
            }
        
        # 融合特征
        if self.fusion_method == "attention":
            # 使用注意力机制融合
            stacked_features = torch.stack(all_features, dim=1)  # (B, num_modules, prompt_dim)
            attended_features, attention_weights = self.fusion_attention(
                stacked_features, stacked_features, stacked_features
            )
            fused_features = attended_features.mean(dim=1)  # (B, prompt_dim)
            fused_prompts['attention_weights'] = attention_weights
            
        elif self.fusion_method == "concat":
            # 简单拼接
            fused_features = torch.cat(all_features, dim=-1)  # (B, num_modules * prompt_dim)
            # 投影到目标维度
            if not hasattr(self, 'concat_projection'):
                self.concat_projection = nn.Linear(
                    len(all_features) * self.prompt_dim, 
                    self.prompt_dim
                ).to(self.device)
            fused_features = self.concat_projection(fused_features)
            
        elif self.fusion_method == "weighted":
            # 加权融合
            weights = F.softmax(self.fusion_weights, dim=0)
            fused_features = sum(w * f for w, f in zip(weights, all_features))
            
        else:
            # 默认：简单平均
            fused_features = torch.stack(all_features, dim=0).mean(dim=0)
        
        # 生成最终提示
        batch_size = fused_features.size(0)
        
        # 点提示
        points = self.final_generator['points'](fused_features)  # (B, 2)
        points = points.unsqueeze(1)  # (B, 1, 2)
        
        # 框提示
        boxes = self.final_generator['boxes'](fused_features)  # (B, 4)
        boxes = boxes.unsqueeze(1)  # (B, 1, 4)
        
        # 掩码提示
        mask_logits = self.final_generator['masks'](fused_features)  # (B, 32*32)
        masks = mask_logits.view(batch_size, 1, 32, 32)  # (B, 1, 32, 32)
        
        # 置信度
        confidence = self.confidence_predictor(fused_features)  # (B, 1)
        
        fused_prompts.update({
            'points': points,
            'boxes': boxes,
            'masks': masks,
            'confidence': confidence,
            'fused_features': fused_features
        })
        
        return fused_prompts
    
    def forward(self, 
                image_features: torch.Tensor,
                context: Optional[Dict[str, Any]] = None) -> Dict[str, torch.Tensor]:
        """
        前向传播
        
        Args:
            image_features: 图像特征 (B, C, H, W)
            context: 上下文信息
            
        Returns:
            提示字典
        """
        # 生成各个模块的提示
        individual_prompts = self.generate_individual_prompts(image_features, context)
        
        # 融合提示
        fused_prompts = self.fuse_prompts(individual_prompts)
        
        # 添加各模块的详细信息
        fused_prompts['individual_prompts'] = individual_prompts
        
        return fused_prompts
    
    def generate_prompts(self, 
                        image: torch.Tensor,
                        metadata: Optional[Dict[str, Any]] = None) -> Dict[str, torch.Tensor]:
        """
        生成提示
        
        Args:
            image: 输入图像 (B, C, H, W)
            metadata: 元数据
            
        Returns:
            提示字典
        """
        # 提取图像特征
        image_features = self.extract_features(image)
        
        # 生成融合提示
        prompts = self.forward(image_features, metadata)
        
        return prompts
    
    def compute_visual_learning_loss(self, 
                                   predictions: Dict[str, torch.Tensor],
                                   targets: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        计算视觉学习损失
        
        Args:
            predictions: 预测结果
            targets: 目标结果
            
        Returns:
            损失值
        """
        total_loss = 0.0
        
        # 基础提示损失
        total_loss += self.compute_loss(predictions, targets)
        
        # 各模块的特定损失
        if 'individual_prompts' in predictions:
            individual_prompts = predictions['individual_prompts']
            
            # 周期性损失
            if 'periodic' in individual_prompts and self.enable_periodic:
                periodic_loss = self.modules['periodic'].compute_periodic_loss(
                    individual_prompts['periodic'],
                    targets.get('periodic', {})
                )
                total_loss += 0.3 * periodic_loss
            
            # 质量感知损失
            if 'quality' in individual_prompts and self.enable_quality:
                quality_loss = self.modules['quality'].compute_quality_loss(
                    individual_prompts['quality'],
                    targets.get('quality', {})
                )
                total_loss += 0.3 * quality_loss
            
            # 记忆引导损失
            if 'memory' in individual_prompts and self.enable_memory:
                memory_loss = self.modules['memory'].compute_memory_loss(
                    individual_prompts['memory'],
                    targets.get('memory', {})
                )
                total_loss += 0.3 * memory_loss
            
            # 自提示损失
            if 'self_prompt' in individual_prompts and self.enable_self_prompt:
                self_prompt_loss = self.modules['self_prompt'].compute_self_prompt_loss(
                    individual_prompts['self_prompt'],
                    targets.get('self_prompt', {})
                )
                total_loss += 0.2 * self_prompt_loss
        
        # 置信度损失
        if 'confidence' in predictions and 'confidence' in targets:
            confidence_loss = F.mse_loss(
                predictions['confidence'],
                targets['confidence']
            )
            total_loss += 0.1 * confidence_loss
        
        return total_loss
    
    def get_module_status(self) -> Dict[str, bool]:
        """获取各模块状态"""
        return {
            'periodic': self.enable_periodic,
            'quality': self.enable_quality,
            'memory': self.enable_memory,
            'self_prompt': self.enable_self_prompt
        }
    
    def enable_module(self, module_name: str, enable: bool = True):
        """启用/禁用特定模块"""
        if module_name == 'periodic':
            self.enable_periodic = enable
        elif module_name == 'quality':
            self.enable_quality = enable
        elif module_name == 'memory':
            self.enable_memory = enable
        elif module_name == 'self_prompt':
            self.enable_self_prompt = enable
        else:
            raise ValueError(f"未知模块: {module_name}")
        
        self.logger.info(f"模块 {module_name} {'启用' if enable else '禁用'}")
    
    def get_detailed_analysis(self, 
                            image: torch.Tensor,
                            context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        获取详细分析结果
        
        Args:
            image: 输入图像 (B, C, H, W)
            context: 上下文信息
            
        Returns:
            详细分析结果
        """
        with torch.no_grad():
            # 生成提示
            prompts = self.generate_prompts(image, context)
            
            analysis = {
                'fused_prompts': {
                    'points': prompts['points'].cpu().numpy(),
                    'boxes': prompts['boxes'].cpu().numpy(),
                    'masks': prompts['masks'].cpu().numpy(),
                    'confidence': prompts['confidence'].cpu().numpy()
                }
            }
            
            # 各模块的详细分析
            if 'individual_prompts' in prompts:
                individual_analysis = {}
                
                if 'periodic' in prompts['individual_prompts']:
                    periodic_prompts = prompts['individual_prompts']['periodic']
                    individual_analysis['periodic'] = {
                        'period_probs': periodic_prompts.get('period_probs', torch.tensor(0)).cpu().numpy(),
                        'phase_info': periodic_prompts.get('phase_info', torch.tensor(0)).cpu().numpy()
                    }
                
                if 'quality' in prompts['individual_prompts']:
                    quality_prompts = prompts['individual_prompts']['quality']
                    individual_analysis['quality'] = {
                        'quality_scores': quality_prompts.get('quality_scores', torch.tensor(0)).cpu().numpy(),
                        'quality_weights': quality_prompts.get('quality_weights', torch.tensor(0)).cpu().numpy()
                    }
                
                if 'memory' in prompts['individual_prompts']:
                    memory_prompts = prompts['individual_prompts']['memory']
                    individual_analysis['memory'] = {
                        'memory_attention_weights': memory_prompts.get('memory_attention_weights', torch.tensor(0)).cpu().numpy(),
                        'retrieved_masks': memory_prompts.get('retrieved_masks', torch.tensor(0)).cpu().numpy()
                    }
                
                analysis['individual_analysis'] = individual_analysis
            
            return analysis
